"""
优化的回测引擎
支持参数扫描和优化
"""
import pandas as pd
import numpy as np
from typing import Dict, List, Optional, Tuple
from datetime import datetime
from src.data.processor import KlineProcessor
from src.strategy.optimized_scalping import OptimizedScalpingStrategy
from src.strategy.parameter_optimizer import (
    StrategyParams, OptimizationResult, ParameterOptimizer
)
from src.trading.position import PositionManager
from src.trading.risk import RiskManager, RiskControl
from src.trading.exit_strategy import ExitStrategy
from src.costs.calculator import CostCalculator
from src.utils.logger import get_logger

logger = get_logger(__name__)


class OptimizedBacktestEngine:
    """优化的回测引擎"""

    def __init__(self, config: Dict, initial_balance: float = 10000):
        """初始化"""
        self.config = config
        self.initial_balance = initial_balance
        self.current_balance = initial_balance

        # 初始化组件
        self.processor = KlineProcessor()
        self.cost_calculator = CostCalculator(
            maker_fee=config.get('costs', {}).get('maker_fee', 0.02),
            taker_fee=config.get('costs', {}).get('taker_fee', 0.05),
            funding_rate=config.get('costs', {}).get('funding_rate', 0.03)
        )
        self.exit_strategy = ExitStrategy(config.get('exit', {}))

        # 风险管理
        risk_control = RiskControl(
            stop_loss_pct=config.get('risk', {}).get('stop_loss', 1.0),
            take_profit_pct=config.get('risk', {}).get('take_profit', 2.0),
            max_positions=config.get('risk', {}).get('max_positions', 3),
            max_daily_loss_pct=config.get('risk', {}).get('max_daily_loss', 5.0),
            max_drawdown_pct=config.get('risk', {}).get('max_drawdown', 10.0),
            max_position_bars=config.get('trading', {}).get('max_position_bars', 20),
        )
        self.risk_manager = RiskManager(risk_control, initial_balance)

        # 回测统计
        self.stats = {
            'bars_processed': 0,
            'trades_executed': 0,
            'peak_balance': initial_balance,
            'lowest_balance': initial_balance,
        }

    def run_backtest_single(self, symbols: List[str],
                           klines_data: Dict[str, Dict[int, pd.DataFrame]],
                           params: StrategyParams,
                           start_date: datetime = None,
                           end_date: datetime = None) -> OptimizationResult:
        """
        运行单个参数组合的回测

        Returns:
            优化结果
        """
        # 重置状态
        self.current_balance = self.initial_balance
        self.stats = {
            'bars_processed': 0,
            'trades_executed': 0,
            'peak_balance': self.initial_balance,
            'lowest_balance': self.initial_balance,
        }

        # 初始化策略
        strategy = OptimizedScalpingStrategy(self.config, params)
        position_manager = PositionManager()

        # 处理K线数据
        for symbol in symbols:
            if symbol not in klines_data:
                continue

            for timeframe, df in klines_data[symbol].items():
                if df is None or df.empty:
                    continue

                # 添加技术指标
                df = self.processor.add_technical_indicators(df, params.to_dict())
                # 直接存储到缓存中，避免重复处理
                key = f"{symbol}_{timeframe}"
                self.processor.cache[key] = df

        # 回测主循环
        for symbol in symbols:
            df_5m = self.processor.get_klines(symbol, 5)
            if df_5m is None or len(df_5m) < 50:
                continue

            self._backtest_symbol_with_params(
                symbol, df_5m, strategy, position_manager, params
            )

        # 生成结果
        result = self._generate_optimization_result(position_manager)
        return result

    def run_backtest_grid(self, symbols: List[str],
                         klines_data: Dict[str, Dict[int, pd.DataFrame]],
                         params_list: List[StrategyParams],
                         start_date: datetime = None,
                         end_date: datetime = None) -> Tuple[ParameterOptimizer, List[OptimizationResult]]:
        """
        运行网格搜索，测试多个参数组合

        Args:
            symbols: 交易对列表
            klines_data: K线数据
            params_list: 参数列表
            start_date: 开始日期
            end_date: 结束日期

        Returns:
            (优化器, 结果列表)
        """
        optimizer = ParameterOptimizer()
        results = []

        total = len(params_list)
        for i, params in enumerate(params_list, 1):
            logger.info(f"[{i}/{total}] 回测参数组合: "
                       f"EMA({params.ema_fast_period}/{params.ema_slow_period}), "
                       f"RSI({params.rsi_period}), "
                       f"SL({params.stop_loss_pct:.2%})")

            try:
                result = self.run_backtest_single(
                    symbols, klines_data, params, start_date, end_date
                )
                optimizer.add_result(result)
                results.append(result)

                logger.info(f"  结果: {result}")

            except Exception as e:
                logger.error(f"  参数组合回测失败: {e}")
                continue

        optimizer.print_summary()
        return optimizer, results

    def _backtest_symbol_with_params(self, symbol: str, df_5m: pd.DataFrame,
                                    strategy: OptimizedScalpingStrategy,
                                    position_manager: PositionManager,
                                    params: StrategyParams):
        """对单个交易对进行回测"""
        logger.debug(f"开始回测: {symbol}, 共{len(df_5m)}条K线")

        # 获取其他时间框架的数据
        df_15m = self.processor.get_klines(symbol, 15)
        df_1h = self.processor.get_klines(symbol, 60)
        df_4h = self.processor.get_klines(symbol, 240)

        # 记录每个持仓的开仓K线索引
        position_entry_bar = {}

        # 逐K线回测
        for i in range(50, len(df_5m)):
            current_bar = df_5m.iloc[i]
            current_time = pd.to_datetime(current_bar['timestamp'], unit='ms')

            # 收集各时间框架数据
            kline_data = {
                5: df_5m.iloc[max(0, i-49):i+1].copy(),
            }

            if df_15m is not None and len(df_15m) > 0:
                kline_data[15] = df_15m.tail(50).copy()

            if df_1h is not None and len(df_1h) > 0:
                kline_data[60] = df_1h.tail(50).copy()

            if df_4h is not None and len(df_4h) > 0:
                kline_data[240] = df_4h.tail(50).copy()

            # 检查现有持仓
            position = position_manager.get_active_position(symbol)

            # 平仓逻辑
            if position:
                # 计算持仓K线数
                bars_held = i - position_entry_bar.get(position.trade_id, i)

                # 使用新的出场策略，将params转为dict以便出场策略读取
                params_dict = {
                    'max_position_bars': params.max_position_bars,
                    'take_profit': params.take_profit_pct * 100,
                    'stop_loss': params.stop_loss_pct * 100,
                    'trailing_stop_pct': 0.5,  # 追踪止损0.5%
                }

                should_exit, reason = self.exit_strategy.should_exit(
                    position, current_bar, kline_data, bars_held, params_dict
                )

                if should_exit:
                    self._close_position(position, current_bar, position_manager, reason)
                    if position.trade_id in position_entry_bar:
                        del position_entry_bar[position.trade_id]
            else:
                # 生成交易信号
                signal, confidence = strategy.generate_signal(kline_data)

                if signal in ["buy", "sell"] and confidence > params.signal_threshold:
                    trade = self._open_position(symbol, signal, current_bar, position_manager,
                                      confidence, i)
                    if trade:
                        position_entry_bar[trade.trade_id] = i

            self.stats['bars_processed'] += 1

    def _open_position(self, symbol: str, signal: str, current_bar: pd.Series,
                      position_manager: PositionManager, confidence: float,
                      bar_index: int):
        """开仓"""
        # 检查风险
        position_size = self.current_balance * 0.05  # 每笔5%的资金

        if not self.risk_manager.can_open_position(position_size):
            return

        price = current_bar['close']

        trade = position_manager.open_position(
            symbol=symbol,
            side="long" if signal == "buy" else "short",
            price=price,
            quantity=position_size / price,
            timestamp=pd.to_datetime(current_bar['timestamp'], unit='ms').to_pydatetime()
        )

        # 记录开仓K线索引
        trade._entry_bar_index = bar_index

        self.stats['trades_executed'] += 1
        logger.debug(f"[开仓] {symbol} {signal.upper()}, 价格={price:.2f}, 信心度={confidence:.2f}")

        return trade

    def _close_position(self, position, current_bar: pd.Series,
                       position_manager: PositionManager, reason: str):
        """平仓"""
        price = current_bar['close']

        # 计算盈亏
        if position.side == "long":
            pnl = (price - position.entry_price) * position.entry_qty
        else:
            pnl = (position.entry_price - price) * position.entry_qty

        # 扣除费用
        cost_info = self.cost_calculator.calculate_profit_loss(
            entry_price=position.entry_price,
            exit_price=price,
            quantity=position.entry_qty,
            trade_type=self.config.get('trading', {}).get('trade_type', 'spot')
        )

        pnl_after_cost = cost_info['net_pnl']

        # 平仓
        closed_trade = position_manager.close_position(
            trade_id=position.trade_id,
            exit_price=price,
            exit_qty=position.entry_qty,
            timestamp=pd.to_datetime(current_bar['timestamp'], unit='ms').to_pydatetime(),
            reason=reason,
            cost_info=cost_info
        )

        # 更新账户余额
        self.current_balance += pnl_after_cost
        self.risk_manager.update_balance(self.current_balance)

        # 更新统计
        self.stats['peak_balance'] = max(self.stats['peak_balance'], self.current_balance)
        self.stats['lowest_balance'] = min(self.stats['lowest_balance'], self.current_balance)

        logger.debug(f"[平仓] {position.symbol} {reason}, 盈亏={pnl_after_cost:.2f}, "
                    f"账户余额={self.current_balance:.2f}")

    def _calculate_pnl_pct(self, side: str, entry_price: float, current_price: float) -> float:
        """计算盈亏百分比"""
        if side == "long":
            return ((current_price - entry_price) / entry_price) * 100
        else:
            return ((entry_price - current_price) / entry_price) * 100

    def _generate_optimization_result(self, position_manager: PositionManager) -> OptimizationResult:
        """生成优化结果"""
        trades = position_manager.get_closed_trades()
        stats = position_manager.calculate_stats()

        total_return = self.current_balance - self.initial_balance
        return_pct = (total_return / self.initial_balance) if self.initial_balance > 0 else 0

        max_drawdown_pct = (
            (self.stats['lowest_balance'] - self.stats['peak_balance']) / self.stats['peak_balance'] * 100
        ) if self.stats['peak_balance'] > 0 else 0

        # 计算夏普比率 (简化版)
        if len(trades) > 0:
            pnls = [t.pnl for t in trades]
            pnl_std = np.std(pnls) if len(pnls) > 1 else 1.0
            sharpe = (np.mean(pnls) / pnl_std * np.sqrt(252)) if pnl_std > 0 else 0
        else:
            sharpe = 0.0

        result = OptimizationResult(
            params=self.current_params if hasattr(self, 'current_params') else StrategyParams(),
            total_return=return_pct,
            win_rate=stats['win_rate'] / 100 if stats['win_rate'] > 0 else 0,
            sharpe_ratio=sharpe,
            max_drawdown=abs(max_drawdown_pct / 100),
            profit_factor=stats['profit_factor'],
            trades_count=stats['total_trades'],
        )

        return result
